import numpy as np
from scipy.interpolate import RegularGridInterpolator
from scipy.spatial import KDTree


class Interpolator:
    def __init__(self, grid_data, property_array, is_regular=True, k=4):
        self.is_regular = is_regular  # 默认为True，代表原网格规则网格；若为False,则代表不确定，由_is_regular_grid函数进行判断。
        self.k = k  # 非规则网格的最近邻点数量，具有默认值
        self.grid_data = grid_data  # 网格坐标点
        self.property_array = property_array  # 属性值数组
        self.x, self.y, self.z = self.unique_no_sort(self.grid_data[:, 0]), self.unique_no_sort(
            self.grid_data[:, 1]), self.unique_no_sort(self.grid_data[:, 2])  # 获取唯一的XYZ坐标

    def _is_regular_grid(self, x_unique, y_unique, z_unique, tolerance=0.5):
        """
        检查 x, y, z 是否分别等间距，允许误差在指定容差内。
        等间距的定义是数组元素间距的最大值和最小值之差不超过容差。
        参数：
            x_unique (array-like): x 方向唯一值的数组。
            y_unique (array-like): y 方向唯一值的数组。
            z_unique (array-like): z 方向唯一值的数组。
            tolerance (float): 允许的最大间距变化范围,具有默认值。
        返回：
            bool: 如果 x, y, z 都等间距返回 True，否则返回 False。
        """
        def is_equally_spaced(arr):
            """
            检查单个数组是否等间距。
            """
            if len(arr) < 2:
                return True
            diffs = np.abs(np.diff(arr))
            return np.max(diffs) - np.min(diffs) <= tolerance

        return (
                is_equally_spaced(x_unique) and
                is_equally_spaced(y_unique) and
                is_equally_spaced(z_unique)
        )

    @staticmethod
    def _weighted_interpolation(distances, values):
        """
        使用距离加权的插值方法
        参数:
            distances (np.ndarray): 距离数组
            values (np.ndarray): 值数组
        返回:
            float: 插值结果
        """
        weights = 1 / (distances + 1e-12)  # 避免除零
        weights /= weights.sum()  # 归一化权重
        return np.dot(weights, values)

    def _interpolate_with_kdtree(self, points, values, query_points):
        """
        使用 KDTree 和加权插值方法进行非规则网格插值
        参数:
            points (np.ndarray): 网格点坐标
            values (np.ndarray): 网格点的值
            query_points (np.ndarray): 查询点坐标
        返回:
            np.ndarray: 插值结果
        """
        tree = KDTree(points)
        distances, indices = tree.query(query_points, k=self.k)
        interpolated_values = np.zeros(len(query_points))
        for i in range(len(query_points)):
            interpolated_values[i] = self._weighted_interpolation(distances[i], values[indices[i]])
        return interpolated_values

    def unique_no_sort(self, array):
        # 获取唯一值和它们在原始数组中的索引   按照原始数组的顺序排序唯一值
        array = np.round(array, decimals=4)
        unique_values, indices = np.unique(array, return_index=True)
        unique_values_in_order = array[np.sort(indices)]
        return unique_values_in_order

    def perform_interpolation(self, new_point):
        """
        执行插值并返回结果。
        参数:
        new_point, 新的插值坐标点。
        """
        new_point = np.round(new_point, decimals=4)
        if self.is_regular:
            # 规则网格
            prop = self.property_array.reshape(len(self.x), len(self.y), len(self.z))
            interpolator = RegularGridInterpolator((self.x, self.y, self.z), prop, method='nearest')
            interpolated_values = interpolator(new_point)
        elif self._is_regular_grid(self.x, self.y, self.z):
            # 规则网格
            prop = self.property_array.reshape(len(self.x), len(self.y), len(self.z))
            interpolator = RegularGridInterpolator((self.x, self.y, self.z), prop, method='nearest')
            interpolated_values = interpolator(new_point)
        else:
            # 非规则网格
            interpolated_values = self._interpolate_with_kdtree(self.grid_data, self.property_array, new_point)
        return interpolated_values.flatten()
